import numpy as np
import matplotlib.pyplot as plt

# 聚类的类
class JuLei:
    def __init__(self, ge_shu=2, zui_da_ci_shu=100, cha_zhi=0.0001):
        self.ge_shu = ge_shu  # 要聚成几类
        self.zui_da_ci_shu = zui_da_ci_shu  # 最多迭代多少次
        self.cha_zhi = cha_zhi  # 中心点变化小于这个就停
        self.zhong_xin = None  # 存放中心点
        self.biao_qian = None  # 每个点的类别标签

    def xun_lian(self, shu_ju):
        # 训练模型，shu_ju是输入的数据
        ge_shu, te_zheng = shu_ju.shape  # 多少个样本，多少个特征

        # 随便选几个点当初始中心点
        suiji_bianhao = np.random.choice(ge_shu, self.ge_shu, replace=False)
        self.zhong_xin = shu_ju[suiji_bianhao]

        for _ in range(self.zui_da_ci_shu):
            # 给每个点分标签
            self.biao_qian = self.fen_lei(shu_ju)

            # 重新算中心点
            xin_zhong_xin = self.suan_zhong_xin(shu_ju)

            # 看看中心点变了多少，差不多就停
            bianhua = 0
            for i in range(self.ge_shu):
                bianhua = bianhua + np.linalg.norm(xin_zhong_xin[i] - self.zhong_xin[i])
            if bianhua < self.cha_zhi:
                break

            self.zhong_xin = xin_zhong_xin

    def fen_lei(self, shu_ju):
        # 给每个数据点分配类别
        biao_qian_list = []
        for dian in shu_ju:
            # 算每个点到所有中心点的距离
            ju_li_list = []
            for zx in self.zhong_xin:
                ju_li = np.sqrt(np.sum((dian - zx) **2))  # 欧氏距离
                ju_li_list.append(ju_li)
            # 找最近的那个中心点的序号
            zui_jin = np.argmin(ju_li_list)
            biao_qian_list.append(zui_jin)
        return np.array(biao_qian_list)

    def suan_zhong_xin(self, shu_ju):
        # 计算新的中心点
        xin_zhong_xin_list = []
        for i in range(self.ge_shu):
            # 找出属于第i类的所有点
            lei_nei_dian = []
            for j in range(len(shu_ju)):
                if self.biao_qian[j] == i:
                    lei_nei_dian.append(shu_ju[j])
            lei_nei_dian = np.array(lei_nei_dian)
            # 算平均值当新中心点
            xin_zx = np.mean(lei_nei_dian, axis=0)
            xin_zhong_xin_list.append(xin_zx)
        return np.array(xin_zhong_xin_list)


# 测试一下
if __name__ == "__main__":
    # 自己造点数据
    def zao_shu_ju():
        np.random.seed(10)
        shu_ju1 = np.random.normal(0, 1, (100, 2))  # 第一类
        shu_ju2 = np.random.normal(5, 1, (100, 2))  # 第二类
        shu_ju3 = np.random.normal(10, 1, (100, 2)) # 第三类
        return np.vstack((shu_ju1, shu_ju2, shu_ju3))

    # 生成数据
    my_shu_ju = zao_shu_ju()

    # 聚类
    my_julei = JuLei(ge_shu=3)  # 聚成3类
    my_julei.xun_lian(my_shu_ju)
    biao_qian = my_julei.biao_qian
    zhong_xin = my_julei.zhong_xin

    # 画图
    plt.scatter(my_shu_ju[:, 0], my_shu_ju[:, 1], c=biao_qian, s=30)
    plt.scatter(zhong_xin[:, 0], zhong_xin[:, 1], c='red', marker='x', s=200)
    plt.title('聚类结果')
